import pandas as pd
import os
import math
from sklearn.manifold import TSNE
from sklearn.decomposition import PCA
from pandas import DataFrame,Series
from scipy.fftpack import fft
import numpy as np
# from __future__ import division, print_function # 引入3.x版本的除法和打印
from matplotlib import pyplot as plt
# 在notebook中显示绘图结果
from scipy.fftpack import fft
from tqdm import tqdm

#时域特征提取
def  psfeatureTime(data,p1,p2):
#     #均值
    df_mean=data[p1:p2].mean()
#     #方差
    df_var=data[p1:p2].var()
#     #标准差
    df_std=data[p1:p2].std()
#     #均方根
    df_rms=math.sqrt(pow(df_mean,2) + pow(df_std,2))
#     #偏度
    df_skew=data[p1:p2].skew()
#     峭度
    df_kurt=data[p1:p2].kurt()
    sum=0
    for i in range(p1,p2):
        sum+=math.sqrt(abs(data[i]))
    #波形因子
    df_boxing=df_rms / (abs(data[p1:p2]).mean())
    #峰值因子
    df_fengzhi=(max(data[p1:p2])) / df_rms
    #脉冲因子
    df_maichong=(max(data[p1:p2])) / (abs(data[p1:p2]).mean())
    #裕度因子
    df_yudu=(max(data[p1:p2])) / pow((sum/(p2-p1)),2)
    featuretime_list = [df_rms,df_skew,df_kurt,df_boxing,df_fengzhi,df_maichong,df_yudu]
#     featuretime_list = [df_kurt,df_boxing,df_fengzhi,df_maichong,df_yudu]
    return featuretime_list

def getFeatureTime(dataPath, flag, col):
    allFileNameList = []
    for root, dirs, files in os.walk(dataPath):
        allFileNameList = files
    fileNameList = []
    fileId = []
    for i in allFileNameList:
        if i[-5:-4] == flag:
            fileNameList.append(i)
            fileId.append(i[0:-6])

    datalist = []

    for name in tqdm(fileNameList):
        # print(os.listdir(dataPath))
        data = pd.read_csv(dataPath + name)
        psdata = psfeatureTime(data[col], 0, 79999)
        datalist.append(psdata)
    timeFeature = pd.DataFrame(datalist, index=fileId)
    return timeFeature

def getFeatureTimeData(dataPath):
    # TIME P  F,  AI1
    TFAI1 = getFeatureTime(dataPath, "F", "ai1")
    TFAI2 = getFeatureTime(dataPath, "F", "ai2")
    TBAI1 = getFeatureTime(dataPath, "B", "ai1")
    TBAI2 = getFeatureTime(dataPath, "B", "ai2")
    TIMEdata = pd.concat([TFAI1, TFAI2, TBAI1, TBAI2], axis=1)
    colnamelist = []
    for i in range(28):
        colnamelist.append("c" + str(i))

    TIMEdata.columns = colnamelist
    return TIMEdata

### 原始数据地址
ppath= "../../xuelangyun_d/Motor_tain/Positive/"
npath= "../../xuelangyun_d/Motor_tain/Negative/"
tpath= "../../xuelangyun_d/Motor_testP/"

TimeDataN=getFeatureTimeData(npath)
# TimeDataN=pd.read_csv("TimeDataN.csv",index_col=0)

TimeDataP=getFeatureTimeData(ppath)
# TimeDataP=pd.read_csv("TimeDataP.csv",index_col=0)

TimeDataTest=getFeatureTimeData(tpath)
# TimeDataTest.to_csv("TimeDataTest.csv")

print("时域提取完毕！")
#频域特征提取
def data2fft(data):
    N = 4000
    fft_y = np.abs(fft(data))  # 数据转为fft
    fft_y = fft_y[0:N]  # 取fft中的前N个频率值
    fft_frequency = np.arange(N)  # 构造频率数据
    fft_y_reduce50 = fft_y[50:].tolist()  # 抛弃频率的前50行数据
    fft_y_max = max(fft_y_reduce50)  # 找出原频率数据中最大的幅值
    fft_frequency_max = fft_y_reduce50.index(max(fft_y_reduce50)) + 50  # 找出原频率数据中幅值最大对应的频率值

    return fft_y, fft_frequency, fft_y_max, fft_frequency_max


def oneFftData(train_data_fft):
    data_fft = {}
    columns = ["F_ai1", "F_ai2", "B_ai1", "B_ai2"]
    for i in columns:
        Fs = 51200  # 采样频率
        flag1 = 1000  # 数据起始位置
        flag2 = 1000 + Fs  # 数据结束位置
        train_data_fft = train_data_fft[flag1:flag2]

        fft_y, fft_frequency, fft_y_max, fft_frequency_max = data2fft(train_data_fft[i])
        data_fft[i + "_fft_y_max"] = fft_y_max
        data_fft[i + "_fft_frequency_max"] = fft_frequency_max

    data_fft["ai1_frequency_max_diff"] = data_fft["F_ai1_fft_frequency_max"] - data_fft["B_ai1_fft_frequency_max"]
    data_fft["ai1_frequency_max_mean"] = (data_fft["F_ai1_fft_frequency_max"] + data_fft["B_ai1_fft_frequency_max"]) / 2
    data_fft["ai2_frequency_max_diff"] = data_fft["F_ai2_fft_frequency_max"] - data_fft["B_ai2_fft_frequency_max"]
    data_fft["ai2_frequency_max_mean"] = (data_fft["F_ai2_fft_frequency_max"] + data_fft["B_ai2_fft_frequency_max"]) / 2
    data_fft["F_frequency_max_diff"] = data_fft["F_ai1_fft_frequency_max"] - data_fft["F_ai2_fft_frequency_max"]
    data_fft["F_frequency_max_mean"] = (data_fft["F_ai1_fft_frequency_max"] + data_fft["F_ai2_fft_frequency_max"]) / 2
    data_fft["B_frequency_max_diff"] = data_fft["B_ai1_fft_frequency_max"] - data_fft["B_ai2_fft_frequency_max"]
    data_fft["B_frequency_max_mean"] = (data_fft["B_ai1_fft_frequency_max"] + data_fft["B_ai2_fft_frequency_max"]) / 2
    return data_fft


def fftDataFrame(dataPath):
    allFileNameList = []
    for root, dirs, files in os.walk(dataPath):
        allFileNameList = files
    fileNameList = []
    fileId = []
    for i in allFileNameList:
        if i[-5:-4] == "F":
            fileNameList.append(i)
            fileId.append(i[0:-6])

    datalist = []
    for name in tqdm(fileId):
        data_F = pd.read_csv(dataPath + name + "_F.csv")
        data_B = pd.read_csv(dataPath + name + "_B.csv")
        train_data = pd.DataFrame()
        train_data["F_ai1"] = data_F["ai1"]
        train_data["F_ai2"] = data_F["ai2"]
        train_data["B_ai1"] = data_B["ai1"]
        train_data["B_ai2"] = data_B["ai2"]
        # train_data
        fftData = pd.DataFrame(oneFftData(train_data), index=[name])
        datalist.append(fftData)
    allFftData = pd.concat(datalist, axis=0)
    return allFftData

fftDataP=fftDataFrame(ppath)
# fftDataP.to_csv("fftDataP.csv")
fftDataN=fftDataFrame(npath)
# fftDataN.to_csv("fftDataN.csv")
fftDataT=fftDataFrame(tpath)
# fftDataT.to_csv("fftDataT.csv")
print("频域提取完毕！")

#加上均值和方差的频域代码
def psfeatureFrequency(dataPath, flag, col):
    allFileNameList = []
    for root, dirs, files in os.walk(dataPath):
        allFileNameList = files
    fileNameList = []
    fileId = []
    for i in allFileNameList:
        if i[-5:-4] == flag:
            fileNameList.append(i)
            fileId.append(i[0:-6])

    datalist = []

    for name in fileNameList:
        data = pd.read_csv(dataPath + name)
        Fs = 51200  # 采样频率
        flag1 = 1000  # 数据起始位置
        flag2 = 1000 + Fs  # 数据结束位置
        train_data_fft = data[flag1:flag2]

        N = 4000
        fft_y = np.abs(fft(train_data_fft[col]))  # 数据转为fft
        fft_y = fft_y[0:N]  # 取fft中的前N个频率值
        fft_frequency = np.arange(N)  # 构造频率数据
        fft_y_reduce50 = fft_y[50:].tolist()  # 抛弃频率的前50行数据
        datalist.append(fft_y_reduce50)

    frequencyDataFrame = pd.DataFrame(datalist, index=fileId)
    return frequencyDataFrame

# 'FA1dataMean','FA2dataMean', 'BA1dataMean', 'BA2dataMean', 'FA1dataVar', 'FA2dataVar','BA1dataVar', 'BA2dataVar'],
FA1pdata=psfeatureFrequency(ppath,"F",'ai1')
FA1ndata=psfeatureFrequency(npath,"F",'ai1')

FA2pdata=psfeatureFrequency(ppath,"F",'ai2')
FA2ndata=psfeatureFrequency(npath,"F",'ai2')

BA1pdata=psfeatureFrequency(ppath,"B",'ai1')
BA1ndata=psfeatureFrequency(npath,"B",'ai1')

BA2pdata=psfeatureFrequency(ppath,"B",'ai2')
BA2ndata=psfeatureFrequency(npath,"B",'ai2')

FA1tdata=psfeatureFrequency(tpath,"F",'ai1')
FA2tdata=psfeatureFrequency(tpath,"F",'ai2')
BA1tdata=psfeatureFrequency(tpath,"B",'ai1')
BA2tdata=psfeatureFrequency(tpath,"B",'ai2')

def getFftDataP2WithMeanAndVar(FA1data,FA2data,BA1data,BA2data,fftData):
    FA1dataMean=FA1data.mean(1)
    FA2dataMean=FA2data.mean(1)
    BA1dataMean=BA1data.mean(1)
    BA2dataMean=BA2data.mean(1)
    FA1dataVar=FA1data.var(1)
    FA2dataVar=FA2data.var(1)
    BA1dataVar=BA1data.var(1)
    BA2dataVar=BA2data.var(1)
    F1=pd.DataFrame(FA1dataMean,columns=['FA1dataMean'])
    F2=pd.DataFrame(FA2dataMean,columns=['FA2dataMean'])
    F3=pd.DataFrame(BA1dataMean,columns=['BA1dataMean'])
    F4=pd.DataFrame(BA2dataMean,columns=['BA2dataMean'])
    F5=pd.DataFrame(FA1dataVar,columns=['FA1dataVar'])
    F6=pd.DataFrame(FA2dataVar,columns=['FA2dataVar'])
    F7=pd.DataFrame(BA1dataVar,columns=['BA1dataVar'])
    F8=pd.DataFrame(BA2dataVar,columns=['BA2dataVar'])
    meanAndVar=pd.concat([F1,F2,F3,F4,F5,F6,F7,F8],axis=1)
    newfftData=pd.concat([fftData,meanAndVar],axis=1)
    return  newfftData

newfftDataP=getFftDataP2WithMeanAndVar(FA1pdata,FA2pdata,BA1pdata,BA2pdata,fftDataP)
# newfftDataP.to_csv("fftDataP21.csv")
newfftDataN=getFftDataP2WithMeanAndVar(FA1ndata,FA2ndata,BA1ndata,BA2ndata,fftDataN)
# newfftDataN.to_csv("fftDataN21.csv")
newfftDataT=getFftDataP2WithMeanAndVar(FA1tdata,FA2tdata,BA1tdata,BA2tdata,fftDataT)
# ewfftDataT.to_csv("fftDataT21.csv")

allPdata=pd.concat([TimeDataP,newfftDataP],axis=1)
allNdata=pd.concat([TimeDataN,newfftDataN],axis=1)
allTdata=pd.concat([TimeDataTest,newfftDataT],axis=1)
allPdata.to_csv("allPdata.csv")
allNdata.to_csv('allNdata.csv')
allTdata.to_csv('allTdata.csv')